#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
使用 libbpf 的 CO-RE (Compile Once – Run Everywhere) 方法实现的 eBPF 捕获服务
"""

import os
import time
import logging
import threading
import socket
import struct
import json
import subprocess
from datetime import datetime
from bson import ObjectId
from db.mongo import get_db

logger = logging.getLogger(__name__)

# 尝试导入 bcc 库，用于检查是否可用
try:
    import bcc
    BCC_AVAILABLE = True
except ImportError:
    BCC_AVAILABLE = False
    logger.warning("BCC库不可用")

# 尝试导入 libbpf 库，用于 CO-RE 方法
try:
    # 检查系统中是否安装了 libbpf 相关库
    result = subprocess.run(['ldconfig', '-p'], capture_output=True, text=True)
    if 'libbpf.so' in result.stdout:
        LIBBPF_AVAILABLE = True
        logger.info("libbpf 库可用")
    else:
        # 检查 bpftool 是否可用
        result = subprocess.run(['which', 'bpftool'], capture_output=True, text=True)
        if result.returncode == 0:
            LIBBPF_AVAILABLE = True
            logger.info("bpftool 可用，路径: %s", result.stdout.strip())
        else:
            LIBBPF_AVAILABLE = False
            logger.warning("libbpf 库和 bpftool 工具都不可用")
except Exception as e:
    LIBBPF_AVAILABLE = False
    logger.warning(f"libbpf 库不可用: {e}")

class LibbpfCaptureService:
    """使用 libbpf 的 eBPF 捕获服务"""

    # 捕获线程
    _capture_thread = None

    # 是否正在运行
    _is_running = False

    # 要捕获的服务列表
    _services = []

    # BPF 程序路径
    _bpf_program_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'bpf', 'http_filter.bpf.o')

    @staticmethod
    def is_available():
        """
        检查 libbpf 是否可用

        Returns:
            bool: 是否可用
        """
        # 再次检查系统中是否安装了 libbpf 相关库
        try:
            # 检查 libbpf 库是否存在
            result = subprocess.run(['ldconfig', '-p'], capture_output=True, text=True)
            if 'libbpf.so' in result.stdout:
                logger.info("libbpf 库可用")
                return True

            # 检查 bpftool 是否可用
            result = subprocess.run(['which', 'bpftool'], capture_output=True, text=True)
            if result.returncode == 0:
                logger.info("bpftool 可用，路径: %s", result.stdout.strip())
                return True

            # 检查 BCC 库是否可用
            if BCC_AVAILABLE:
                logger.info("BCC 库可用，将使用 BCC 实现 eBPF 功能")
                return True

            logger.warning("libbpf 库、bpftool 工具和 BCC 库都不可用")
            return False
        except Exception as e:
            logger.warning(f"libbpf 库不可用: {e}")
            return False

    @staticmethod
    def start_capture(interface='eth0', services=None):
        """
        启动捕获

        Args:
            interface: 网络接口
            services: 要捕获的服务列表，格式为 [{"service_id": "xxx", "service_name": "xxx", "service_ip": "xxx", "service_port": xxx}]

        Returns:
            bool: 是否成功启动
        """
        if LibbpfCaptureService._is_running:
            logger.warning("捕获服务已经在运行")
            return False

        # 保存服务列表
        LibbpfCaptureService._services = services or []

        # 启动捕获线程
        LibbpfCaptureService._is_running = True
        LibbpfCaptureService._capture_thread = threading.Thread(
            target=LibbpfCaptureService._capture_worker,
            args=(interface, LibbpfCaptureService._services),
            daemon=True
        )
        LibbpfCaptureService._capture_thread.start()

        logger.info(f"已启动 libbpf 捕获服务，网络接口: {interface}，服务数量: {len(LibbpfCaptureService._services)}")
        return True

    @staticmethod
    def stop_capture():
        """
        停止捕获

        Returns:
            bool: 是否成功停止
        """
        if not LibbpfCaptureService._is_running:
            logger.warning("捕获服务未在运行")
            return False

        # 停止捕获线程
        LibbpfCaptureService._is_running = False
        if LibbpfCaptureService._capture_thread:
            LibbpfCaptureService._capture_thread.join(timeout=5)
            LibbpfCaptureService._capture_thread = None

        logger.info("已停止 libbpf 捕获服务")
        return True

    @staticmethod
    def is_running():
        """
        检查是否正在运行

        Returns:
            bool: 是否正在运行
        """
        return LibbpfCaptureService._is_running

    @staticmethod
    def get_services():
        """
        获取正在捕获的服务列表

        Returns:
            list: 服务列表
        """
        return LibbpfCaptureService._services

    @staticmethod
    def _capture_worker(interface, services):
        """
        捕获工作线程

        Args:
            interface: 网络接口
            services: 要捕获的服务列表
        """
        # 检查 libbpf 是否可用
        if not LibbpfCaptureService.is_available():
            logger.warning("libbpf 库不可用，无法启动捕获服务")
            return False

        logger.info(f"开始使用 libbpf (CO-RE) 捕获网络接口 {interface} 上的流量")
        service_info = []
        for s in services:
            name = s.get('service_name', 'unknown')
            ip = s.get('service_ip', '0.0.0.0')
            port = s.get('service_port', 0)
            service_info.append(f"{name}({ip}:{port})")
        logger.info(f"监控服务列表: {service_info}")

        try:
            # 编译 BPF 程序（如果需要）
            if not os.path.exists(LibbpfCaptureService._bpf_program_path):
                LibbpfCaptureService._compile_bpf_program()

            # 使用 bpftool 加载 BPF 程序
            cmd = [
                "bpftool", "prog", "load", LibbpfCaptureService._bpf_program_path,
                "/sys/fs/bpf/http_filter", "type", "socket"
            ]
            subprocess.run(cmd, check=True)

            # 创建原始套接字
            sock = socket.socket(socket.AF_PACKET, socket.SOCK_RAW, socket.htons(0x0003))
            sock.bind((interface, 0))

            # 获取 BPF 程序文件描述符
            cmd = ["bpftool", "prog", "show", "name", "http_filter", "-j"]
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            prog_info = json.loads(result.stdout)
            prog_fd = prog_info[0]["id"]

            # 附加 BPF 程序到套接字
            cmd = ["bpftool", "prog", "attach", str(prog_fd), "msg_verdict", "dev", interface]
            subprocess.run(cmd, check=True)

            # 打开 perf 缓冲区
            cmd = ["bpftool", "map", "show", "name", "http_events", "-j"]
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            map_info = json.loads(result.stdout)
            map_fd = map_info[0]["id"]

            # 创建 perf 事件
            cmd = ["perf", "record", "-e", f"bpf-output:map_fd={map_fd}", "-a"]
            perf_process = subprocess.Popen(cmd)

            # 添加服务到 BPF 映射
            for service in services:
                ip = service.get('service_ip', '0.0.0.0')
                port = service.get('service_port', 0)
                ip_int = struct.unpack("!I", socket.inet_aton(ip))[0]

                # 添加到 BPF 映射
                cmd = [
                    "bpftool", "map", "update", "name", "services",
                    "key", f"hex {ip_int:08x} {port:04x}",
                    "value", "hex 01"
                ]
                subprocess.run(cmd, check=True)

            # 等待停止信号
            while LibbpfCaptureService._is_running:
                time.sleep(1)

            # 清理
            perf_process.terminate()
            sock.close()

            # 分离 BPF 程序
            cmd = ["bpftool", "prog", "detach", str(prog_fd), "msg_verdict", "dev", interface]
            subprocess.run(cmd, check=True)

            # 卸载 BPF 程序
            cmd = ["rm", "-f", "/sys/fs/bpf/http_filter"]
            subprocess.run(cmd, check=True)

            logger.info("libbpf 捕获线程已退出")

        except Exception as e:
            logger.error(f"libbpf 捕获线程异常: {e}")
            LibbpfCaptureService._is_running = False

    @staticmethod
    def _compile_bpf_program():
        """编译 BPF 程序"""
        try:
            # 获取 BPF 程序源文件路径
            src_path = LibbpfCaptureService._bpf_program_path.replace('.o', '.c')

            # 检查必要的头文件是否存在
            if not os.path.exists("/usr/include/linux/bpf.h"):
                logger.error("缺少必要的头文件: /usr/include/linux/bpf.h")
                return False

            if not os.path.exists("/usr/include/asm/types.h") and not os.path.exists("/usr/include/linux/types.h"):
                logger.error("缺少必要的头文件: /usr/include/asm/types.h 或 /usr/include/linux/types.h")
                return False

            # 使用 clang 编译 BPF 程序
            cmd = [
                "clang", "-O2", "-target", "bpf", "-c", src_path,
                "-o", LibbpfCaptureService._bpf_program_path
            ]

            # 添加调试信息
            logger.info(f"执行命令: {' '.join(cmd)}")
            result = subprocess.run(cmd, capture_output=True, text=True)
            if result.returncode != 0:
                logger.error(f"编译 BPF 程序失败: {result.stderr}")
                return False

            logger.info(f"BPF 程序编译成功: {LibbpfCaptureService._bpf_program_path}")

        except Exception as e:
            logger.error(f"编译 BPF 程序失败: {e}")
            raise

    @staticmethod
    def _process_event(event_data):
        """
        处理 BPF 事件

        Args:
            event_data: 事件数据
        """
        try:
            # 解析事件数据
            ts = struct.unpack("Q", event_data[0:8])[0]
            pid = struct.unpack("I", event_data[8:12])[0]
            saddr = struct.unpack("I", event_data[12:16])[0]
            daddr = struct.unpack("I", event_data[16:20])[0]
            sport = struct.unpack("H", event_data[20:22])[0]
            dport = struct.unpack("H", event_data[22:24])[0]
            direction = struct.unpack("B", event_data[24:25])[0]
            data = event_data[25:].decode('utf-8', errors='ignore')

            # 转换 IP 地址
            saddr_str = socket.inet_ntoa(struct.pack("I", saddr))
            daddr_str = socket.inet_ntoa(struct.pack("I", daddr))

            # 检查是否匹配服务
            if not LibbpfCaptureService._filter_service(saddr_str, daddr_str, sport, dport, LibbpfCaptureService._services):
                return

            # 处理 HTTP 请求或响应
            if direction == 1:  # 请求
                LibbpfCaptureService._process_http_request(saddr_str, daddr_str, sport, dport, data, LibbpfCaptureService._services)
            else:  # 响应
                LibbpfCaptureService._process_http_response(saddr_str, daddr_str, sport, dport, data, LibbpfCaptureService._services)

        except Exception as e:
            logger.error(f"处理 BPF 事件异常: {e}")

    @staticmethod
    def _filter_service(saddr, daddr, sport, dport, services):
        """
        过滤服务

        Args:
            saddr: 源IP
            daddr: 目标IP
            sport: 源端口
            dport: 目标端口
            services: 服务列表

        Returns:
            bool: 是否匹配
        """
        for service in services:
            service_ip = service.get('service_ip')
            service_port = service.get('service_port')

            # 检查是否匹配
            if (saddr == service_ip and sport == service_port) or (daddr == service_ip and dport == service_port):
                return True

        return False

    @staticmethod
    def _process_http_request(saddr, daddr, sport, dport, data, services):
        """
        处理 HTTP 请求

        Args:
            saddr: 源IP
            daddr: 目标IP
            sport: 源端口
            dport: 目标端口
            data: HTTP 数据
            services: 服务列表
        """
        try:
            # 解析 HTTP 请求
            lines = data.split('\r\n')
            if not lines:
                return

            # 解析请求行
            request_line = lines[0].split(' ')
            if len(request_line) < 3:
                return

            method = request_line[0]
            path = request_line[1]
            version = request_line[2]

            # 解析请求头
            headers = {}
            for i in range(1, len(lines)):
                line = lines[i]
                if not line:
                    break

                parts = line.split(':', 1)
                if len(parts) == 2:
                    headers[parts[0].strip()] = parts[1].strip()

            # 解析请求体
            body = None
            if 'Content-Length' in headers:
                content_length = int(headers['Content-Length'])
                if content_length > 0:
                    body_start = data.find('\r\n\r\n') + 4
                    if body_start > 0 and body_start < len(data):
                        body = data[body_start:body_start + content_length]

            # 构建请求 URL
            host = headers.get('Host', f"{daddr}:{dport}")
            url = f"http://{host}{path}"

            # 查找匹配的服务
            service_id = None
            service_name = None
            for service in services:
                if service.get("service_ip") == daddr and service.get("service_port") == dport:
                    service_id = service.get("service_id")
                    service_name = service.get("service_name")
                    break

            if not service_id:
                return

            # 保存请求信息
            LibbpfCaptureService._save_request(service_id, service_name, method, url, headers, body)

        except Exception as e:
            logger.error(f"处理 HTTP 请求异常: {e}")

    @staticmethod
    def _process_http_response(saddr, daddr, sport, dport, data, services):
        """
        处理 HTTP 响应

        Args:
            saddr: 源IP
            daddr: 目标IP
            sport: 源端口
            dport: 目标端口
            data: HTTP 数据
            services: 服务列表
        """
        try:
            # 解析 HTTP 响应
            lines = data.split('\r\n')
            if not lines:
                return

            # 解析状态行
            status_line = lines[0].split(' ', 2)
            if len(status_line) < 3:
                return

            version = status_line[0]
            status_code = int(status_line[1])
            status_message = status_line[2] if len(status_line) > 2 else ""

            # 解析响应头
            headers = {}
            for i in range(1, len(lines)):
                line = lines[i]
                if not line:
                    break

                parts = line.split(':', 1)
                if len(parts) == 2:
                    headers[parts[0].strip()] = parts[1].strip()

            # 解析响应体
            body = None
            if 'Content-Length' in headers:
                content_length = int(headers['Content-Length'])
                if content_length > 0:
                    body_start = data.find('\r\n\r\n') + 4
                    if body_start > 0 and body_start < len(data):
                        body = data[body_start:body_start + content_length]

            # 查找匹配的服务
            service_id = None
            service_name = None
            for service in services:
                if service.get("service_ip") == saddr and service.get("service_port") == sport:
                    service_id = service.get("service_id")
                    service_name = service.get("service_name")
                    break

            if not service_id:
                return

            # 保存响应信息
            LibbpfCaptureService._save_response(service_id, service_name, status_code, headers, body)

        except Exception as e:
            logger.error(f"处理 HTTP 响应异常: {e}")

    @staticmethod
    def _save_request(service_id, service_name, method, url, headers, body):
        """
        保存请求信息

        Args:
            service_id: 服务ID
            service_name: 服务名称
            method: 请求方法
            url: 请求URL
            headers: 请求头
            body: 请求体
        """
        try:
            # 创建临时请求记录
            db = get_db()
            db.temp_http_requests.insert_one({
                "service_id": service_id,
                "service_name": service_name,
                "method": method,
                "url": url,
                "headers": headers,
                "body": body,
                "created_at": datetime.now()
            })

            logger.debug(f"已保存请求: {method} {url}")

        except Exception as e:
            logger.error(f"保存请求信息异常: {e}")

    @staticmethod
    def _save_response(service_id, service_name, status_code, headers, body):
        """
        保存响应信息

        Args:
            service_id: 服务ID
            service_name: 服务名称
            status_code: 状态码
            headers: 响应头
            body: 响应体
        """
        try:
            # 创建临时响应记录
            db = get_db()
            db.temp_http_responses.insert_one({
                "service_id": service_id,
                "service_name": service_name,
                "status_code": status_code,
                "headers": headers,
                "body": body,
                "created_at": datetime.now()
            })

            logger.debug(f"已保存响应: {status_code} from {service_name}")

        except Exception as e:
            logger.error(f"保存响应信息异常: {e}")


